#ifndef __THREAD_POOL_HPP__
#define __THREAD_POOL_HPP__

#include <functional>
#include <vector>
#include <memory>
#include <thread>
#include <queue>
#include <mutex>
#include <atomic>
#include <condition_variable>
#include <cassert>
#include <future>
#include <string>
#include <unordered_map>
#include <iostream>

namespace thread_pool
{
    namespace thread_infomation
    {
        class ThreadInfomation
        {
        private:
            std::thread::id _threadId;
            int _threadNumber;    // This corresponds to the thread set index
            bool _isSingleFinish; // Whether a single thread ends.
            std::mutex _mutex;

        public:
            ThreadInfomation(const std::thread::id &threadId = std::thread::id(), const int &threadNumber = -1, const bool &isSingleFinish = false)
                : _threadId(threadId), _threadNumber(threadNumber), _isSingleFinish(isSingleFinish) {}

            ~ThreadInfomation() = default;

            ThreadInfomation &operator=(const ThreadInfomation &other)
            {
                if (this != &other)
                {
                    _threadId = other._threadId;
                    _threadNumber = other._threadNumber;
                    _isSingleFinish = other._isSingleFinish;
                }
                return *this;
            }

            void setIsSingleFinish()
            {
                std::unique_lock<std::mutex> lock(_mutex);
                _isSingleFinish = true;
            }

            void info()
            {
                std::unique_lock<std::mutex> lock(_mutex);
                printf("[threadId:%u]->[threadNumber:%d]->[isFinish:%s] : task done!!!\n", _threadId, _threadNumber, isSingleFinish().c_str());
            }

        private:
            std::string isSingleFinish()
            {
                return _isSingleFinish ? "yes" : "no";
            }

        }; // class ThreadInfomation

    } // namespace thread_pool_infomation

    namespace task_pool
    {
        template <class Function>
        class TaskPool
        {
        private:
            std::queue<Function> _tasks; // task set.
            std::mutex _mutex;           //
        public:
            bool push(const Function &task)
            {
                std::unique_lock<std::mutex> lock(_mutex); // Unique lock ensures the uniqueness of the lock.
                _tasks.push(task);
                return true;
            }

            bool pop(Function &task)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                if (_tasks.empty())
                {
                    return false;
                }
                task = _tasks.front();
                _tasks.pop();
                return true;
            }

            bool empty() const
            {
                std::unique_lock<std::mutex> lock(_mutex);
                return _tasks.empty();
            }

        }; // class task_pool

    } // namespace task_pool

    class ThreadPool
    {
    private:
        typedef std::function<void(int threadId)> task_type; // The threadId parameter is the subscript of the _threads member.

    private:
        std::vector<std::unique_ptr<std::thread>> _threads;                        // thread set. The subscript is similar to the thread id.
        std::vector<std::shared_ptr<std::atomic<bool>>> _isSingleFinish;           // Whether a single thread ends. No end is false, end is true.
        task_pool::TaskPool<task_type *> _taskPool;                                // task pool.
        std::atomic<bool> _isStop;                                                 // Whether the thread pool is stopped. If it is true, it does stop. If it is false, it does not stop.
        std::atomic<bool> _isEnd;                                                  // Whether the thread pool ends.  No end is false, end is true.
        std::atomic<int> _waitCount;                                               // Thread wait count.
        std::mutex _mutex;                                                         // Thread pool mutex.
        std::condition_variable _isWait;                                           // condition judgment. If it is true, it does not wait. If it is flase, it waits.
        std::unordered_map<int, thread_infomation::ThreadInfomation> _threadInfos; // thread info set

    public: // Copy prohibition
        ThreadPool(const ThreadPool &) = delete;
        ThreadPool(ThreadPool &&) = delete;
        ThreadPool &operator=(const ThreadPool &) = delete;
        ThreadPool &operator=(ThreadPool &&) = delete;

    public:
        int size() const { return static_cast<int>(_threads.size()); }
        int waitCount() const { return static_cast<int>(_waitCount); }
        std::thread &getThread(int index) const { return *(_threads[index]); }

    public:
        explicit ThreadPool() { initialize(); }

        ThreadPool(int thread_count)
        {
            assert(thread_count >= 0);
            initialize();
            resize(thread_count);
#ifdef DEBUG_RUN_INFO
            std::cout << "TreadPool : constructor done!" << std::endl;
#endif
        }

        ~ThreadPool() { stop(true); }

        void resize(int thread_count)
        {
            assert(thread_count >= 0);
            if (!_isEnd && !_isStop)
            {
                int old_thread_count = static_cast<int>(_threads.size());
                if (old_thread_count == thread_count)
                {
                    return;
                }
                if (old_thread_count < thread_count) // dilatation
                {
                    _threads.resize(thread_count);
                    _isSingleFinish.resize(thread_count);

                    for (int i = old_thread_count; i < thread_count; ++i)
                    {
                        _isSingleFinish[i] = std::make_shared<std::atomic<bool>>(false);
                        std::cout << i;
                        setThreads(i);
                    }
                }
                else
                {
                    for (int i = old_thread_count - 1; i >= thread_count; --i)
                    {
                        *(_isSingleFinish[i]) = true;
                        _threads[i]->detach(); // The thread ends when the main function of the thread terminates
#ifdef DEBUG_THREAD_INFO
                        _threadInfos[i].setIsSingleFinish();
                        _threadInfos[i].info();
#endif
                    }
                    {
                        // stop the detached threads that were waiting
                        std::unique_lock<std::mutex> lock(_mutex);
                        _isWait.notify_all();
                    }
                    _threads.resize(thread_count);
                    _isSingleFinish.resize(thread_count);
                }
            }
        }

        std::function<void(int)> pull()
        {
            task_type *result = nullptr;
            _taskPool.pop(result);
            task_type _result;
            if (result)
            {
                _result = *result;
            }
            return _result;
        }

        template <typename Function, typename... Args>
        auto submit(Function &&function, Args &&...args) -> std::future<decltype(function(0, args...))>
        {
            // To bind user functions and parameters, just pass in the first parameter
            auto packaging_task = std::make_shared<std::packaged_task<decltype(function(0, args...))(int)>>(
                std::bind(std::forward<Function>(function), std::placeholders::_1, std::forward<Args>(args)...));
            // The binding function is further encapsulated and is new
            auto task = new std::function<void(int threadId)>([packaging_task](int threadId)
                                                              { (*packaging_task)(threadId); });
            _taskPool.push(task);
            {
                std::unique_lock<std::mutex> lock(_mutex);
                _isWait.notify_one();
            }
#ifdef DEBUG_RUN_INFO
            std::cout << "sumbit done" << std::endl;
#endif
            // Returns a future object, contains the result of the user function.
            return packaging_task->get_future();
        }

        template <typename Function>
        auto submit(Function &&function) -> std::future<decltype(function(0))>
        {
            // To bind user functions and parameters, just pass in the first parameter
            auto packaging_task = std::make_shared<std::packaged_task<decltype(function(0))(int)>>(std::forward<Function>(function));
            // The binding function is further encapsulated and is new
            auto task = new std::function<void(int threadId)>([packaging_task](int threadId)
                                                              { (*packaging_task)(threadId); });
            _taskPool.push(task);
            {
                std::unique_lock<std::mutex> lock(_mutex);
                _isWait.notify_one();
            }
#ifdef DEBUG_RUN_INFO
            std::cout << "submit done" << std::endl;
#endif
            // Returns a future object, contains the result of the user function.
            return packaging_task->get_future();
        }

    private:
        void initialize()
        {
            _isStop = false;
            _isEnd = false;
            _waitCount = 0;
        }

        void setThreads(int index)
        {
            std::shared_ptr<std::atomic<bool>> __isSingleFinish(_isSingleFinish[index]);
            auto task_function = [this, index, __isSingleFinish]()
            {
                std::atomic<bool> &___isSingleFinish = *__isSingleFinish;
                task_type *task;
                bool isPop = _taskPool.pop(task);
                while (true)
                {
                    while (isPop)
                    {
                        (*task)(index);
#ifdef DEBUG_THREAD_INFO // After the task is completed, the thread information is output
                        _threadInfos[index].info();
#endif
                        if (___isSingleFinish) // the thread is wanted to stop, return even if the queue is not empty yet
                        {
                            return;
                        }
                        else
                        {
                            isPop = _taskPool.pop(task);
                        }
                    }
                    std::unique_lock<std::mutex> lock(_mutex);
                    ++_waitCount;
                    _isWait.wait(lock, [this, &task, &isPop, &___isSingleFinish]()
                                 { isPop = _taskPool.pop(task);
                                    return isPop || _isEnd || ___isSingleFinish; });
                    --_waitCount;
                    {
                        return;
                    }
                }
            };
            _threads[index].reset(new std::thread(task_function));
#ifdef DEBUG_RUN_INFO
            std::cout << "reset done" << std::endl;
#endif
#ifdef DEBUG_THREAD_INFO // Thread start record information
            _threadInfos[index] = thread_infomation::ThreadInfomation(_threads[index]->get_id(), index, (*__isSingleFinish));
#endif
        }

        void stop(bool isWait = false)
        {
            if (!isWait)
            {
                if (_isStop)
                {
                    return;
                }
                _isStop = true;
                for (int i = 0, n = size(); i < n; ++i)
                {
                    *(_isSingleFinish[i]) = true;
                }
                clearTaskPool();
            }
            else
            {
                if (_isStop || _isEnd)
                {
                    return;
                }
                _isEnd = true;
            }
            {
                // stop all waiting threads
                std::unique_lock<std::mutex> lock(_mutex);
                _isWait.notify_all();
            }
            for (int i = 0, n = size(); i < n; ++i)
            {
                if (_threads[i]->joinable())
                {
                    _threads[i]->join();
                }
            }
            clearTaskPool();
            _threads.clear();
            _isSingleFinish.clear();
            _threadInfos.clear();
        }

        void clearTaskPool()
        {
            task_type *task;
            while (_taskPool.pop(task))
            {
                delete task;
            }
        }

    }; // class ThreadPool

} // namespace thread_pool

#endif // #ifndef __THREAD_POOL_HPP__
